[DSA] Make the fused top-k exact on an overflowing threshold bin - #37941
xiaofei-zheng wants to merge 1 commit into
Conversation
8a746f3 to
3dbc304
Compare
ae5db33 to
7a83d7f
Compare
|
@HaiShaw could you add the The gate currently exits with For context: this is a correctness fix in the fused DSA indexer top-k, found while profiling our own GLM-5.2 decode runs on MI355X. When a threshold coarse bin holds more than |
|
@BBuf @yuan-luo @DarkSharpness please review |
|
/tag-and-rerun-ci |
DarkSharpness
left a comment
There was a problem hiding this comment.
Some comments on performance
| // Early-out for a bit-identical candidate set: the radix passes cannot | ||
| // separate it and would only re-derive what phase 3 already staged. One | ||
| // distinct exact key means the candidates are interchangeable, so any | ||
| // kMaxNumTie-subset is correct, including the one already in tie.values. | ||
| // Comparing min against max of the key is exact for any value, and costs | ||
| // the one scan that non-degenerate overflow rows pay on top. | ||
| { | ||
| uint32_t key_min = 0xFFFFFFFFu; | ||
| uint32_t key_max = 0u; | ||
| for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) { | ||
| if (val >= v_lo && val < v_hi) { | ||
| const auto key = extract_exact_bin(val); | ||
| key_min = min(key_min, key); | ||
| key_max = max(key_max, key); | ||
| } | ||
| }); | ||
| // Reduce per-thread extrema: two atomics per thread, not per candidate. | ||
| if (tx == 0) { | ||
| handle->histogram[0][0] = 0xFFFFFFFFu; | ||
| handle->histogram[0][1] = 0u; | ||
| } | ||
| __syncthreads(); | ||
| atomicMin(&handle->histogram[0][0], key_min); | ||
| atomicMax(&handle->histogram[0][1], key_max); | ||
| __syncthreads(); | ||
| const bool bit_identical = handle->histogram[0][0] == handle->histogram[0][1]; | ||
| __syncthreads(); // all threads read the scratch before the loop clears it | ||
| if (bit_identical) { | ||
| // equal_count > kMaxNumTie on entry, so phase 3 filled the whole buffer. | ||
| const auto above_count = smem->count_gt; | ||
| const auto remain_topk = above_count < topk ? topk - above_count : 0; | ||
| handle_tie(smem->tie.values, problem, above_count, kMaxNumTie, remain_topk, handle); | ||
| return; | ||
| } | ||
| } |
There was a problem hiding this comment.
Why do we need this pass? It's not common that all keys are equal in fp32. We should optimize for hot path. This actually adds to additional linear scan cost in normal pass.
| for (uint32_t round = 0; round < 4 && cand_count > kMaxNumTie && remain > 0; ++round) { | ||
| const uint32_t shift = 24 - round * 8; | ||
|
|
||
| if (tx < kRadixSize) handle->histogram[0][tx] = 0; | ||
| __syncthreads(); | ||
| for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) { | ||
| if (val >= v_lo && val < v_hi) { | ||
| const auto key = extract_exact_bin(val); | ||
| if ((key & mask) == prefix) atomicAdd(&handle->histogram[0][(key >> shift) & 0xFFu], 1); | ||
| } | ||
| }); | ||
| __syncthreads(); | ||
|
|
||
| refine_find_threshold(cand_count, remain, handle); | ||
| const auto match = handle->match; | ||
|
|
||
| for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) { | ||
| if (val >= v_lo && val < v_hi) { | ||
| const auto key = extract_exact_bin(val); | ||
| if ((key & mask) == prefix && ((key >> shift) & 0xFFu) > match.bin) { | ||
| const auto pos = atomicAdd(&smem->count_gt, 1); | ||
| if (pos < topk) [[likely]] | ||
| problem.emit(pos, idx); | ||
| } | ||
| } | ||
| }); | ||
|
|
||
| prefix |= match.bin << shift; | ||
| mask |= 0xFFu << shift; | ||
| remain -= match.above_count; | ||
| cand_count = match.equal_count; | ||
| __syncthreads(); // `match` is read above and overwritten by the next pass | ||
| } |
There was a problem hiding this comment.
Is this efficient enough? This would involve 8 linear pass. Please reduce that to at most 4 passes.
| if (tx == 0) smem->count_eq = 0; | ||
| __syncthreads(); | ||
| for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) { | ||
| if (val >= v_lo && val < v_hi) { | ||
| const auto key = extract_exact_bin(val); | ||
| if ((key & mask) == prefix) { | ||
| const auto slot = atomicAdd(&smem->count_eq, 1); | ||
| if (slot < kMaxNumTie) smem->tie.values[slot] = {val, idx}; | ||
| } | ||
| } | ||
| }); |
There was a problem hiding this comment.
same, this passes should be eliminated i guess
DarkSharpness
left a comment
There was a problem hiding this comment.
Some comments on performance
7e0bc53 to
c05fe16
Compare
|
@DarkSharpness thanks — all three addressed in c05fe16, and measured on MI355X. The min/max early-out is gone. You are right that it was paying a scan on every refinement to serve a degenerate case. Without it the radix rounds still terminate correctly on an all-equal row: they consume all 32 bits and then truncate a bit-identical set, which is exact. Removing it also freed the tie staging buffer, which this path does not need — it re-derives the candidates rather than trusting what the collect pass staged there — so the buffer now carries the histogram. Down to at most four passes. Two changes got there. The radix is 12 bits instead of 8, using that freed buffer (it is exactly The staging pass is eliminated too, fused into the last round's emit the same way. Kernel time, 200 iterations after warmup:
The hot path is unchanged, as it should be — it still only pays the On correctness after the rewrite: the suite is 288/288 on this branch. Run against the unfixed kernel on main it is 8 failed / 2 passed, the two passes being the |
c05fe16 to
fcd3905
Compare
|
@DarkSharpness when you have a moment — the performance changes you asked for are in The branch has also been rebased onto latest
The remaining red checks are unrelated to this diff, as far as I can tell:
Happy to dig into any of those if you would rather they were green first. |
fcd3905 to
9a57b69
Compare
|
@HaiShaw thanks for the earlier approval. Could you help nudge this one along, or advise on how to proceed? @DarkSharpness requested changes on Sep 9 with three performance comments. All three are addressed in Every CI job that exercises this kernel is now green on both vendors, and |
|
Hi. Could you please temporarily hold on until #38798 is merged? Sorry we're actively working on that and it should land within next few days. I will check this PR in detail later. This topk kernel is very sensitive in register usage on CUDA, and any seemingly irrelevant test might lead to wild performance data, so I will take a very detailed look at this PR later. Please stay tuned. Thanks :) |
…verflowing threshold bin The coarse histogram bin is derived from the top bits of the fp16 cast of each score, so one bin spans about a quarter-binade. Phase 3 stages threshold-bin candidates into the fixed kMaxNumTie = 2048 staging buffer and phase 4 clips tie_count = min(count_eq, kMaxNumTie), so once a threshold bin holds more than 2048 elements handle_tie ranks whichever arrived first and the kernel returns a wrong selected-value multiset with no diagnostic. This is the shared DSA indexer top-k used by DeepSeek-V4.1 decode, prefill and the DeepGEMM candidate path, and it fails silently: every emitted index is a valid KV position, so nothing downstream detects the wrong selection. Refine the overflowing bin on the exact key instead of truncating it: up to three radix rounds (12/12/8 bits) over the order-preserving key from extract_exact_bin. It exits once the refined set fits the staging buffer, or once all 32 bits are consumed -- the key is injective on fp32 bit patterns, so the survivors are bit-identical by then and truncating them is legitimate. The fast path branches on count_eq > kMaxNumTie and is otherwise unchanged; kMaxNumTie, kBlockSize and the shared-memory footprint are untouched (the refinement histogram overlays the dead tie staging buffer in a union). This is a correctness fix, not a speedup: the hot path is unchanged. The CUDA-only TopKCluster path still truncates; refining there needs cluster-wide rather than block-local histogram and emit counters, and is left for a follow-up. Ported from sgl-project#37941, which does not apply cleanly after sgl-project#38829/sgl-project#39098 refactored this file. Identifier adaptations: smem->tie.handle -> smem->tie_handle, smem->tie.values -> smem->tie_values, smem->tie.refine_hist -> smem->refine_hist; the local sites name the counters count_gt/count_eq rather than above_count/equal_count. The PR's ROCm __launch_bounds__ hunk is intentionally dropped (CUDA-only deployment). Verified on H20 (SM90, cc 9.0): - The new test fails 8/10 on the unpatched kernel (the two all_equal controls pass, pinning the fault to truncation) and passes 10/10 patched. - test/registered/kernels/ops/attention/test_topk_v2.py: 298 passed. - cuobjdump resource usage and sizeof(Smem) are byte-identical to baseline: 31-32 registers, zero spills, shared memory unchanged.
|
Sounds good, we'll hold. Two things worth flagging while we wait. The truncation this PR fixes is unchanged in #38798 — the And please don't try to carry this change into #38798; it is safer for us to re-port it afterwards. The refinement re-reads candidates from global memory and relies on On your register point: we only have MI355X here, so every number in this PR is gfx950 (57 to 78 VGPRs, no spills, LDS unchanged). We cannot measure the CUDA side at all. If that is the main concern, it would help if someone with an NVIDIA box could check it — or tell us which shapes and metric you want and we will get them into the PR another way. Ping me when #38798 lands and I'll rebase, re-measure and re-request review. |
The coarse bin comes from the top bits of the fp16 cast, so one bin spans a quarter-binade. The collect pass stages threshold-bin candidates under `if (pos < kMaxNumTie)` and phase 4 clips `min(count_eq, kMaxNumTie)`, so an overflowing bin leaves `handle_tie` ranking an arrival-order subset and the kernel returns a wrong selected-value multiset with no diagnostic. Found while profiling GLM-5.2 decode on MI355X: device-side counters over an agentic serving run enter that path on ~0.04% of rows across all 8 ranks, with a largest observed bin of 4139 against the 2048 cap and none of those rows bit-identical. sgl-project#35257 reports a related problem from the CUDA side. sgl-project#39648 restructured this kernel but left all four truncation sites unchanged. Refine the overflowing bin on the exact key instead: up to three radix rounds over the order-preserving key from `extract_exact_bin`, 12 / 12 / 8 bits wide, exiting once the refined set fits the staging buffer or once all 32 bits are consumed, where the survivors are bit-identical and truncating them is exact. Each round emits its own "above" set while building the next round's histogram, so the refinement costs at most four input passes. The histogram overlays `tie_values`, which this path re-derives rather than trusting, so LDS is unchanged at 26928 B. The fast path branches on `count_eq > kMaxNumTie` and is otherwise untouched. Also pin the ROCm `__launch_bounds__` occupancy argument to the physical wave floor: HIP reads it as waves per SIMD where CUDA reads it as blocks per SM, so `kOccupancy` asked ROCm for 2 waves/SIMD, below the 4 a 1024-thread block already forces. The CUDA spelling stays because `topk_small_batch_cluster_kernel` shadows `kOccupancy` with a template parameter. Testing. Pre-existing cases draw from `torch.randn` and tolerate MAX_PERMIT_ERROR = 5, so they never reach the overflow. Added narrow, narrow_bin, tiny and two_values distributions that collapse a row into one coarse bin, plus overflow with DUAL_OUTPUT and overflow on an unaligned ragged window. On the kernel on main those 14 cases fail while the 2 all_equal controls pass, which pins the fault to truncation rather than to the binning. All 304 pass here. MI355X: the hot path is unchanged (5.324 vs 5.484 us at 6x8192, 12.094 vs 12.267 at 4x32768); rows that overflow now pay for work previously skipped by returning a wrong answer (narrow_bin 17.051 -> 21.684, two_values 36.194 -> 118.811 at 4x32768). 57 -> 79 VGPRs, zero spills. Known gap: the CUDA-only `TopKCluster` path still truncates. Refining there needs cluster-wide rather than block-local histogram and emit counters. Co-authored-by: Cursor <cursoragent@cursor.com>
9a57b69 to
e503afa
Compare
|
Rebased onto #39648 and re-ported, now A correction to what I said earlier: I claimed the padding change would break the refinement's counting invariant. That was wrong. Also worth flagging: #39648 restructured the kernel but left all four truncation sites unchanged, so the defect is exactly as it was. Verified on MI355X — 304 pass here, and against the kernel on main the 14 new overflow cases fail while the 2 |
Motivation
The coarse histogram bin comes from the top bits of the fp16 cast of each score, so one bin spans a quarter-binade. Phase 3 stages threshold-bin candidates under
if (count_eq < kMaxNumTie)and phase 4 clipstie_count = min(equal_count, kMaxNumTie), so once a threshold bin holds more than 2048 elementshandle_tieranks whichever arrived first and the kernel returns a wrong selected-value multiset with no diagnostic. This is the shared DSA indexer top-k, serving every model routed throughsrt/layers/attention/dsa.Found while profiling GLM-5.2 decode on MI355X: device-side counters over an agentic serving run enter the truncation path on ~0.04% of rows across all 8 ranks (556–598 of 1.1–1.5M per rank), largest threshold bin 4139 against the 2048 cap, none of those rows bit-identical. #35257 reports the same defect independently from the CUDA side — referenced, not closed.
Modifications
Refine the overflowing bin on the exact key instead of truncating it: up to four 8-bit radix passes over the order-preserving key from
extract_exact_bin. It exits once the refined set fits the staging buffer, or once all 32 bits are consumed — the key is injective on fp32 bit patterns, so the survivors are bit-identical by then and truncating them is legitimate. A min/max reduction skips the refinement when the set is already bit-identical.The fast path branches on
equal_count > kMaxNumTieand is otherwise unchanged;kMaxNumTie,kBlockSizeand the shared-memory footprint are untouched.Folded in from #37942 (closed): the ROCm
__launch_bounds__occupancy argument now states the physical wave floor, since HIP reads it as waves per SIMD where CUDA reads it as blocks per SM. Codegen-neutral alone; it matters because the plausible wrong reading caps the allocator at 64 VGPRs and would spill at the 77 this change needs.Accuracy Tests
Pre-existing cases draw from
torch.randnand tolerateMAX_PERMIT_ERROR = 5, so they never reach the overflow and pass on the truncating kernel. This adds narrow, narrow_bin and tiny distributions that collapse a row into one coarse bin, on a register-path (6×8192) and a streaming-path (4×32768) shape, asserting exactly. Against the unfixed kernel:all_equalcontrolsall_equaloverflows too, but bit-identically, so any subset is correct and it must keep passing — that asymmetry pins the fault to truncation rather than mis-binning. All 286 pass with this change; a separate 50-case gate moves 29/50 to 50/50.Benchmarking and Profiling
MI355X (gfx950), weighted J over the captured GLM-5.2 decode shapes, graph-replayed:
A paired decode A/B measures −0.50% / −0.27% step time against a 0.04–0.08% within-arm spread, i.e. parity. This is a correctness fix, not a speedup. Register usage goes 57 → 78 VGPRs with zero spills, and the LDS footprint is unchanged at 27392 B.
Per review feedback the refinement was cut from up to ten input passes to at most four (one to seed the first histogram, then one per round, each fusing the previous round's emit). Kernel time on MI355X, 200 iterations after warmup:
randn(no overflow, hot path)randn(no overflow, hot path)narrow_bin(every row refines)narrow_bin(every row refines)two_values(worst case, 32 bits)two_values(worst case, 32 bits)The hot path is untouched; only rows that actually overflow reach the refinement.
Known gap
The CUDA-only
TopKClusterpath still truncates; refining there needs cluster-wide rather than block-local histogram and emit counters, so it is left for a follow-up. #35257 quotes only the register and streaming sites, so that third one is unreported. Related but distinct, #36807 covers the same class of truncation in the AOTtopk.cu.Checklist
Known gap
The CUDA-only
TopKClusterpath still truncates. Refining there is a different routine: the candidate set is split acrosskClusterSizeranks, so each radix pass needs a cluster-wide histogram all-reduce and a cluster-wide emit counter rather than the block-local ones used here. Left for a follow-up rather than grown into this PR — happy to open a tracking issue.CI States
Latest PR Test (Base): ⏳ Run #35081117275
Latest PR Test (Extra): ❌ Run #35081117023
Latest PR Test (AMD ROCm 10): ❌ Run #35081117212